######## Image Object Detection Using Tensorflow-trained Classifier #########
#
# Author: Evan Juras
# Date: 1/15/18
# Description:
# This program uses a TensorFlow-trained classifier to perform object detection.
# It loads the classifier uses it to perform object detection on an image.
# It draws boxes and scores around the objects of interest in the image.

## Some of the code is copied from Google's example at
## https://github.com/tensorflow/models/blob/master/research/object_detection/object_detection_tutorial.ipynb

## and some is copied from Dat Tran's example at
## https://github.com/datitran/object_detector_app/blob/master/object_detection_app.py

## but I changed it to make it more understandable to me.

# Import packages
import os
import numpy as np
import time
import tensorflow as tf
import sys
import PIL
from pylab import *

# Import utilites
from utils import label_map_util
from utils import visualization_utils as vis_util

NUM_CLASSES = 1
BATCHSIZE = 1
MODEL_NAME = '/home/leo/Downloads/self-trained/ssd512_2k/inference_graph'
ROOT_IMG_PATH = '/home/leo/Downloads/datas/road_car2'
ROOT_IMG_PATH_GEN = '/home/leo/Downloads/datas/road_car2_gen'
PATH_TO_CKPT = os.path.join(MODEL_NAME, 'frozen_inference_graph.pb')
PATH_TO_LABELS = os.path.join('labelmap', 'labelmap.pbtxt')


def save_batch_img(batch_boxes, batch_scores, batch_classes, batch_num, batch_images,batch_filenames,ROOT_IMG_PATH_GEN):
    for (boxes, scores, classes, num, image, filename) in zip(batch_boxes, batch_scores, batch_classes, batch_num, batch_images,batch_filenames):
        vis_util.visualize_boxes_and_labels_on_image_array(
            image,
            np.squeeze(boxes),
            np.squeeze(classes).astype(np.int32),
            np.squeeze(scores),
            category_index,
            use_normalized_coordinates=True,
            line_thickness=4,
            min_score_thresh=0.30)
        # imshow(image)
        # plt.show()
        save_file = os.path.join(ROOT_IMG_PATH_GEN,os.path.basename(filename))
        im = PIL.Image.fromarray(np.uint8(image))
        # print("saving to {}".format(save_file))
        im.save(save_file)


def get_batch_img(BATCHSIZE,loading_idx,all_idx,img_files):
    image_expanded = None
    initial = False
    batch_images = []
    batch_filenames = []
    for i in (range(BATCHSIZE)):
        if (i+loading_idx) >= all_idx:
            break
        # print("opening {}".format(img_files[i+loading_idx]))
        if os.path.exists(img_files[i+loading_idx]) is False:
            print("file not exists {} ".format(img_files[i+loading_idx]))
        image = PIL.Image.open(img_files[i+loading_idx])
        image = array(image)
        batch_images.append(image)
        batch_filenames.append(img_files[i+loading_idx])
        image_loaded = np.expand_dims(image, axis=0)
        if initial is False:
            image_expanded = image_loaded
            initial = True
        else:
            image_expanded = np.concatenate((image_expanded, image_loaded), 0)
    # print("image batch shape:{}".format(image_expanded.shape))
    return image_expanded , batch_images , batch_filenames

if os.path.exists(ROOT_IMG_PATH_GEN) is False:
    os.makedirs(ROOT_IMG_PATH_GEN)

img_files = []
for dir in os.listdir(ROOT_IMG_PATH):
    sub_dir = os.path.join(ROOT_IMG_PATH, dir)
    if os.path.isdir(sub_dir) == True:
        for file in os.listdir(sub_dir):
            file_path = os.path.join(sub_dir, file)
            if os.path.splitext(file_path)[1] == ".jpg":
                img_files.append(file_path)

label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES,
                                                            use_display_name=True)
category_index = label_map_util.create_category_index(categories)

# Load the Tensorflow model into memory.
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

    # cpu_num = int(os.environ.get('CPU_NUM', 3))
    # config = tf.ConfigProto(device_count={"CPU": cpu_num},
    #                         inter_op_parallelism_threads=cpu_num,
    #                         intra_op_parallelism_threads=cpu_num,
    #                         log_device_placement=True)
    # sess = tf.Session(graph=detection_graph,config=config)
    sess = tf.Session(graph=detection_graph)

    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

    # Output tensors are the detection boxes, scores, and classes
    # Each box represents a part of the image where a particular object was detected
    detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')

    # Each score represents level of confidence for each of the objects.
    # The score is shown on the result image, together with the class label.
    detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
    detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')

    # Number of objects detected
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')

    loading_idx = 0
    all_idx = len(img_files)
    loop_time = math.ceil(all_idx / BATCHSIZE)

    print(" all file:{} batch size:{} loop time:{}".format(all_idx,BATCHSIZE,loop_time))
    print(" saving path: {}".format(ROOT_IMG_PATH_GEN))
    for loop in range(loop_time):
        sys.stdout.write('\r>> processing  %d/%d ' % (
            loading_idx, all_idx))
        sys.stdout.flush()

        image_expanded, batch_images, batch_filenames = get_batch_img(BATCHSIZE,loading_idx,all_idx,img_files)

        # start = time.time()
        (batch_boxes, batch_scores, batch_classes, batch_num) = sess.run(
            [detection_boxes, detection_scores, detection_classes, num_detections],
            feed_dict={image_tensor: image_expanded})
        # print("using time : {}".format(time.time() - start))
        save_batch_img(batch_boxes, batch_scores, batch_classes, batch_num, batch_images,batch_filenames,ROOT_IMG_PATH_GEN)
        loading_idx+=BATCHSIZE